
import argparse
import os
import json
import cv2
import numpy as np
from annotation_stat import open_spire_annotations


def main():
    parser = argparse.ArgumentParser(description="Statistic on spire annotations")
    parser.add_argument(
        "--image-dir",
        default="D:\\images_from_video",
        help="path to image dir",
    )
    parser.add_argument(
        "--output-dir",
        default="D:\\",
        help="path to output spire dir"
    )
    args = parser.parse_args()

    spire_img_dir = os.path.join(args.output_dir, 'scaled_images')
    spire_ann_dir = os.path.join(args.output_dir, 'annotations')
    if not os.path.exists(spire_img_dir):
        os.makedirs(spire_img_dir)
    if not os.path.exists(spire_ann_dir):
        os.makedirs(spire_ann_dir)

    input_names = os.listdir(args.image_dir)
    for img_name in input_names:

        if img_name.endswith('.jpg'):
            img_fn = os.path.join(args.image_dir, img_name)
            img = cv2.imread(img_fn)

            open(os.path.join(spire_img_dir, img_name), 'wb').write(
                open(img_fn, 'rb').read())

            spire_dict = {}
            spire_dict['file_name'] = img_name
            spire_dict['height'], spire_dict['width'] = img.shape[0], img.shape[1]
            spire_dict['annos'] = []

            fp = open(os.path.join(spire_ann_dir, spire_dict['file_name']+'.json'), 'w')
            json.dump(spire_dict, fp)


if __name__ == '__main__':
    main()
